#!/usr/bin/env python
#-*- coding: utf-8 -*-

# Author: Edward A. Roualdes
#         University of Kentucky

from __future__ import print_function
import sys
from multiprocessing import Pool
from Bio import SeqIO
from sortBarcode import sortBarcode, sortInds
from sortLoci import sortLoci
from sanity import sane
import os.path
import utils
import hap
import structure
import align as a
import phase
import gc

def init(argv=None):
    if argv is None:
        argv = sys.argv

    print("\nWelcome to NextAllele")
    configfile = utils.parseCommandLine(argv)
    opts = sane(configfile)
    return opts

def main(opts = None):

    if opts['sort']:
        utils.mkdir(
            opts['OUTPUTDIRECTORY'], 
            True,
            opt = ("Please specify a different "
                   "OUTPUTDIRECTORY in the CONFIGFILE "))
        if opts['BARCODEFILE']:
            srt, c = sortBarcode(
                opts['INPUTFILE'], 
                opts['BARCODEFILE'], 
                opts['trim'])
        else:
            srt, c = sortInds(opts['INPUTFILE'])

        srt = sortLoci(srt, opts['BLASTREPORT'], c)
        # write sorted seqs to file
        utils.writeSeqs(srt, opts['OUTPUTDIRECTORY'])
        gc.collect()

    if opts['align']:
        print("Creating primary alignments...")
        p = Pool(processes = opts['cpu'], maxtasksperchild=20)
        p.map(alignTool, a.files2align(opts['OUTPUTDIRECTORY']))
        p.close()
        p.join()

    if (opts['structure'] or opts['phase']): 
        # TODO develop name to phase haplotypes.
        # until then, we can't create structure/PHASE files without re-phasing haplotypes
        # phase haplotypes
        hapDir = hap.callHaps(opts['OUTPUTDIRECTORY'], approx = opts['unphase'])
        # align haplotypes
        p = Pool(processes = opts['cpu'], maxtasksperchild=20)
        p.map(alignTool, a.files2align(hapDir))
        p.close()
        p.join()

    if opts['structure']:
        structure.mkStruct(opts['OUTPUTDIRECTORY'], opts['cleansnp'])

    if opts['phase']:
        phase.mkPhase(opts['OUTPUTDIRECTORY'])

    if opts['runphase']:
        global phaseDir 
        phaseDir = os.path.join(opts['OUTPUTDIRECTORY'], 'phase')
        print('Running PHASE...')
        # run program PHASE on all created files
        p = Pool(processes = opts['cpu'], maxtasksperchild=20)
        p.map(phaseTool, phase.files2phase(phaseDir))
        p.close()
        p.join()
            
    print("NextAllele finished.")

# appease multiprocessing (and GIL)
def alignTool(x):
    return a.aligner(x, opts["ALIGNTOOL"][0], opts["ALIGNTOOL"][1])

def phaseTool(x):
    return phase.runPhase(x, opts['PHASETOOL'], phaseDir)

if __name__ == '__main__':
    opts = init()
    sys.exit(main(opts))
